#include "packet_builder.h"
#include "logger.h"
#include <sstream>
#include "stegit_math.h"
#include "ui_socket.h"

#define HEADER_LENGTH_BIT 1
#define HEADER_MAX_LENGTH_FIELD 16 // two bytes
#define HEADER_MAX_CRC_FIELD 8
#define MIN_PAYLOAD_LENGTH 1
#define MAX_PAYLOAD_LENGTH 255

using namespace std;

const std::string CPacketBuilder::sIniSection = "PACKET_BUILDER";

CPacketBuilder::CPacketBuilder(CPropertiesReader* pConfig,
		IPacketSource* pPacketSource)
		:mError(true),
		mConfig(pConfig),
		mPacketSource(pPacketSource),
		mMaxPayload(0),
		mIgnoreWrongCRC(false),
		mSendPayloadInfo(NPacketSource::PAYLOAD_NO_DATA),
		mSendState(SEND_STATE_WAIT_FOR_NEXT),
		mSendPacket(), // size = 0
		mSendBitCount(0),
		mSendBitIndex(0),
		mSendPayloadLength(0),
		mRecvState(RECV_STATE_WAIT_FOR_NEXT),
		mRecvPacket(),
		mRecvBitCount(0),
		mRecvBitIndex(0),
		mRecvPayloadLength(0)
{
	if (!mConfig) {
		printl(LL_ERR, __FILE__, "No config");
		return;
	}
	string lStrValue = "";

	if (!mConfig->containsKey(sIniSection, "MAX_PAYLOAD")
	|| !mConfig->getValue(sIniSection, "MAX_PAYLOAD", lStrValue)) {
		printl(LL_ERR, __FILE__, "Config error: MAX_PAYLOAD [%s]\n", sIniSection.c_str());
		return;
	}

	int lMaxPayload(0);
	{
		stringstream ss;
		ss << lStrValue;
		ss >> lMaxPayload;
	}
	if (lMaxPayload <= 0) {
		printl(LL_ERR, __FILE__, "Max payload must be > 0.\n");
		return;
	}
	if (lMaxPayload > MAX_PAYLOAD_LENGTH) {
		printl(LL_WARN, __FILE__, "Max payload will be set to %u.\n", MAX_PAYLOAD_LENGTH);
		lMaxPayload = MAX_PAYLOAD_LENGTH;
	}

	mMaxPayload = lMaxPayload;

	printl(0, __FILE__, "max payload: %u\n", lMaxPayload);
	mSendPacket.resize(
			HEADER_LENGTH_BIT +
			HEADER_MAX_LENGTH_FIELD +
			mMaxPayload * 8 +
			HEADER_MAX_CRC_FIELD);
	mRecvPacket.resize(mSendPacket.size());

	if (mConfig->containsKey(sIniSection, "IGNORE_CRC")
	&& mConfig->getValue(sIniSection, "IGNORE_CRC", lStrValue)
	&& lStrValue == "true") {
		mIgnoreWrongCRC = true;
	}

	mError = false;
}

CPacketBuilder::~CPacketBuilder()
{
}

IExtBitSource::EBitValue CPacketBuilder::getBit()
{
	if (mSendState == SEND_STATE_WAIT_FOR_NEXT) {
		// check the old state for EOF
		if (mSendPayloadInfo == NPacketSource::PAYLOAD_EOF) {
			return BIT_EOF;
		}

		mSendPayloadInfo = createNextSendPacket();
		switch (mSendPayloadInfo) {
			case NPacketSource::PAYLOAD_ERROR:
				return BIT_EOF;
			case NPacketSource::PAYLOAD_EOF:
				if (mSendPayloadLength == 0) {
					return BIT_EOF;
				}
			case NPacketSource::PAYLOAD_DATA:
				mSendState = SEND_STATE_NEXT_IS_READY;
				break;
			case NPacketSource::PAYLOAD_NO_DATA:
				// mSendState doesn't change --> continue waiting
				break;
		}
	}

	switch (mSendState) {
		case SEND_STATE_WAIT_FOR_NEXT:
			return BIT_NONE;
		case SEND_STATE_NEXT_IS_READY:
			mSendState = SEND_STATE_PACKET;
			return BIT_NEW_PREAMBLE;
		case SEND_STATE_PACKET:
		{
			EBitValue lBit(mSendPacket[mSendBitIndex] ? BIT_TRUE : BIT_FALSE);
			mSendBitIndex++;
			if (mSendBitIndex >= mSendBitCount) {
				mSendState = SEND_STATE_WAIT_FOR_NEXT;
				CUISocket::getInstance()->sendPacketSendLength(mSendPayloadLength);
			}
			return lBit;
		}
	}
	printl(LL_WARN, __FILE__, "Impossible send state.\n");
	return BIT_EOF;
}

void CPacketBuilder::putBit(bool pBitValue, bool pFirstBit)
{
	static unsigned int lCount(0);
	if(pFirstBit) {
		//printl(0, __FILE__, "First bit after preamble. Last bits: %u\n", lCount);
		lCount = 0;
	}
	lCount++;

	if (pFirstBit) {
		if (mRecvState != RECV_STATE_WAIT_FOR_NEXT) {
			printl(LL_WARN, __FILE__, "Got corrupted packet.");
		}
		mRecvState = RECV_STATE_BEGIN;
		mRecvBitIndex = 0;
	}

	if (mRecvState == RECV_STATE_WAIT_FOR_NEXT) {
		return;
	}

	if (mRecvBitIndex >= mRecvPacket.size()) {
		printl(LL_ERR, __FILE__, "Out of range.");
		mRecvState = RECV_STATE_WAIT_FOR_NEXT;
		mRecvPayloadLength = 0;
		return;
	}
	mRecvPacket[mRecvBitIndex] = pBitValue;
	mRecvBitIndex++;

	switch (mRecvState) {
		case RECV_STATE_BEGIN:
			// copied the length flag
			mRecvState = RECV_STATE_LENGTH_FIELD;
			break;
		case RECV_STATE_LENGTH_FIELD:
			if (mRecvPacket[0] == false) {
				if (mRecvBitIndex == 9) { // one byte length field
					unsigned char lLength;
					readByteFromBoolVector(mRecvPacket, 1, lLength);
					mRecvPayloadLength = lLength;
					mRecvState = RECV_STATE_PAYLOAD;
				}
			}
			else {
				if (mRecvBitIndex == 17) { // two bytes (big endian) length field
					uint16_t lLength;
					readBytesFromBoolVector(mRecvPacket, 1, &lLength, 2);
					mRecvPayloadLength = ntohs(lLength);
					//printl(0, __FILE__, "mRecvPayloadLength (2 bytes) = %u", mRecvPayloadLength);
					mRecvState = RECV_STATE_PAYLOAD;
				}
			}
			break;
		case RECV_STATE_PAYLOAD:
			{
				unsigned int lPayloadStart(mRecvPacket[0] ? 17 : 9);
				unsigned int lPayloadEnd(lPayloadStart + (mRecvPayloadLength * 8) + 8/* for CRC */);

				if (mRecvBitIndex == lPayloadEnd) {
					unsigned char* lPayload = new unsigned char [mRecvPayloadLength];
					readBytesFromBoolVector(mRecvPacket, lPayloadStart, lPayload, mRecvPayloadLength);

					unsigned char lCalcCRC(CStegitMath::CRC8_BlockChecksum(lPayload, mRecvPayloadLength));
					unsigned char lRecvCRC;
					readByteFromBoolVector(mRecvPacket, lPayloadEnd - 8, lRecvCRC);

					bool lCrcError(false);
					if (lCalcCRC == lRecvCRC) {
						//printl(0, __FILE__, "mRecvPayloadLength = %u", mRecvPayloadLength);
						mPacketSource->putPacket(lPayload, mRecvPayloadLength);
					}
					else {
						lCrcError = true;

						printl(LL_ERR, __FILE__, "Wrong CRC sum.");
						if (mIgnoreWrongCRC) {
							mPacketSource->putPacket(lPayload, mRecvPayloadLength);
						}
					}

					CUISocket::getInstance()->sendPacketRecvLength(mRecvPayloadLength, lCrcError);

					delete lPayload;
					mRecvState = RECV_STATE_WAIT_FOR_NEXT;
					mRecvPayloadLength = 0;
				}
			}
			break;
		default:
			printl(LL_ERR, __FILE__, "Wrong receive state.");
			break;
	}
}

unsigned int CPacketBuilder::getMaxFrameLength()
{
	unsigned int lBits(0);

	lBits += 1; // length field flag
	lBits += (mMaxPayload > 255) ? 16 : 8; // length field
	lBits += (mMaxPayload * 8); // max. payload
	lBits += 8; // CRC

	return lBits * 5;
}

NPacketSource::EPayloadInfo CPacketBuilder::createNextSendPacket()
{
	if (mError) {
		return NPacketSource::PAYLOAD_ERROR;
	}

	mSendBitCount = 0;
	mSendPayloadLength = 0;

	unsigned int lLength(0);
	void* lPayload;
	NPacketSource::EPayloadInfo lState;

	// get packet from packet source
	lState = mPacketSource->getPacket(lPayload, lLength, mMaxPayload);

	if (lState == NPacketSource::PAYLOAD_ERROR || !lLength) {
		return lState;
	}

	if (lLength > mMaxPayload) {
		printl(LL_ERR, __FILE__, "Too long payload.\n");
		return NPacketSource::PAYLOAD_ERROR;
	}
	if (lLength < MIN_PAYLOAD_LENGTH || lLength > mMaxPayload) {
		printl(LL_ERR, __FILE__, "Currently only a payload length between %d and %d bytes allowed.",
				MIN_PAYLOAD_LENGTH, mMaxPayload);
		return NPacketSource::PAYLOAD_ERROR;
	}

	unsigned int lBitIndex(0);

	mSendPacket[lBitIndex] = false; // length bit = 0 -> 1 byte for length field
	lBitIndex++;

	// length field
	writeByteToBoolVector(mSendPacket, lBitIndex, lLength);
	lBitIndex += 8;

	// payload
	for (unsigned int i = 0; i < lLength; i++) {
		writeByteToBoolVector(mSendPacket, lBitIndex, static_cast<unsigned char*>(lPayload)[i]);
		lBitIndex += 8;
	}

	// crc summe
	unsigned char lCRC(CStegitMath::CRC8_BlockChecksum(static_cast<unsigned char*>(lPayload), lLength));
	writeByteToBoolVector(mSendPacket, lBitIndex, lCRC);
	lBitIndex += 8;

	mSendBitCount = lBitIndex;
	mSendBitIndex = 0;
	mSendPayloadLength = lLength;

	return lState;
}

bool CPacketBuilder::writeByteToBoolVector
(
		std::vector<bool> &pVector,
		unsigned int pStartIndex,
		unsigned char pByteValue
)
{
	if (pStartIndex + 8 >= pVector.size()) {
		printl(LL_ERR, __FILE__, "writeByteToBoolVector failed.\n");
		return false;
	}

	// MSB is the first bit (lowest vector index) and LSB the last one
	// (highest vector index)
	for (int i = 0; i < 8; i++) {
		pVector[pStartIndex + i] = (pByteValue & 0x80) ? true : false;
		pByteValue <<= 1;
	}
	return true;
}

bool CPacketBuilder::readByteFromBoolVector
(
		std::vector<bool> &pVector,
		unsigned int pStartIndex,
		unsigned char& pByteValue
)
{
	if (pStartIndex + 8 >= pVector.size()) {
		printl(LL_ERR, __FILE__, "readByteFromBoolVector failed.\n");
		return false;
	}

	// MSB is the first bit (lowest vector index) and LSB the last one
	// (highest vector index)
	pByteValue = 0x00;
	for (int i = 0; i < 8; i++) {
		pByteValue <<= 1;
		if (pVector[pStartIndex + i]) {
			pByteValue |= 0x01;
		}
	}
	return true;
}

unsigned int CPacketBuilder::readBytesFromBoolVector
(
		std::vector<bool> &pVector,
		unsigned int pStartIndex,
		void* pBytes,
		unsigned int pLength)
{
	unsigned char* lCurrentByte(static_cast<unsigned char*>(pBytes));

	for (unsigned int i = 0; i < pLength; i++) {
		if (!readByteFromBoolVector(pVector, pStartIndex + (i * 8), *lCurrentByte)) {
			return i;
		}
		lCurrentByte++;
	}
	
	return pLength;
}

